/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2019-2021.
 * Description: support squashfs lzma feature
 * Author: gaoyongliang <gaoyongliang@huawei.com>
 * Create: 2019-01-11
 */

#include <linux/mutex.h>
#include <linux/buffer_head.h>
#include <linux/slab.h>
#include <linux/vmalloc.h>
#include <asm/unaligned.h>

#include "squashfs_fs.h"
#include "squashfs_fs_sb.h"
#include "squashfs.h"
#include "decompressor.h"
#include "page_actor.h"

#ifdef CONFIG_SQUASHFS_LZMA_IMPROVED
#include "lzma_improved.h"
#else
#include <linux/decompress/unlzma.h>
#endif

#ifdef CONFIG_RTOS_EXTEND_SQUASHFS
#include "squashfs_ex.h"
#endif

struct squashfs_lzma {
	void	*input;
	void	*output;
};

#ifdef CONFIG_SQUASHFS_LZMA_DEFAULT
/* decompress_unlzma.c doesn't provide any context in its callbacks... */
static int lzma_error;

static void error(char *m)
{
	ERROR("unlzma error: %s\n", m);
	lzma_error = 1;
}
#endif

static void *lzma_init(struct squashfs_sb_info *msblk, void *buff)
{
	int block_size = max_t(int, msblk->block_size, SQUASHFS_METADATA_SIZE);
	struct squashfs_lzma *stream = kzalloc(sizeof(*stream), GFP_KERNEL);

	if (stream == NULL)
		goto failed;

#ifdef CONFIG_SQUASHFS_CHOICE_DECOMP_BY_MOUNT
	if (msblk->thread_ops == &squashfs_decompressor_percpu) {
		ERROR("-o threads=percpu shouldn't be used while mounting image which is compress by lzma");
		kfree(stream);
		return ERR_PTR(-EINVAL);
	}
#endif

#ifdef CONFIG_RTOS_EXTEND_SQUASHFS
	if (msblk->ex_sb_info) {
		/* It is indeed ex-squashfs image. */
		block_size = max_t(int, block_size, msblk->ex_sb_info->block_size);
	}
#endif
	stream->input = vmalloc(block_size);
	if (stream->input == NULL)
		goto failed;
	stream->output = vmalloc(block_size);
	if (stream->output == NULL)
		goto failed2;

	return stream;

failed2:
	vfree(stream->input);
failed:
	ERROR("Failed to allocate lzma workspace\n");
	kfree(stream);
	return ERR_PTR(-ENOMEM);
}


static void lzma_free(void *strm)
{
	struct squashfs_lzma *stream = strm;

	if (stream) {
		vfree(stream->input);
		vfree(stream->output);
	}
	kfree(stream);
}

#ifdef CONFIG_RTOS_EXTEND_SQUASHFS
static int dup_bio_pages(struct bio *bio, void *buff,
	int bytes, int offset)
{
	struct bvec_iter_all iter_all = {};
	struct bio_vec *bvec = bvec_init_iter_all(&iter_all);
	void *data;
	int copied = 0;
	int avail;

	while (bio_next_segment(bio, &iter_all)) {
		avail = min(bytes, ((int)bvec->bv_len) - offset);
		data = page_address(bvec->bv_page) + bvec->bv_offset;
		memcpy(buff, data + offset, avail);
		copied += avail;
		buff += avail;
		bytes -= avail;
		offset = 0;
	}

	return copied;
}
#endif

static int lzma_uncompress(struct squashfs_sb_info *msblk, void *strm,
	struct bio *bio, int offset, int length,
	struct squashfs_page_actor *output)
{
#ifdef CONFIG_RTOS_EXTEND_SQUASHFS
	struct squashfs_lzma *stream = strm;
	void *buff = stream->input, *data;
	int bytes = length, res;
	struct bio_list *bio_list;
	struct bio *iter;
	int copied;

	if (bio->bi_private) {
		bio_list = bio->bi_private;
		bio_list_for_each(iter, bio_list) {
			copied = dup_bio_pages(iter, buff, bytes, offset);
			buff += copied;
			bytes -= copied;
			offset = 0;
		}
	} else {
		dup_bio_pages(bio, buff, bytes, offset);
	}
#else
	struct bvec_iter_all iter_all = {};
	struct bio_vec *bvec = bvec_init_iter_all(&iter_all);
	struct squashfs_lzma *stream = strm;
	void *buff = stream->input, *data;
	int bytes = length, res;

	while (bio_next_segment(bio, &iter_all)) {
		int avail = min(bytes, ((int)bvec->bv_len) - offset);
		data = page_address(bvec->bv_page) + bvec->bv_offset;
		memcpy(buff, data + offset, avail);
		buff += avail;
		bytes -= avail;
		offset = 0;
	}
#endif

#ifdef CONFIG_SQUASHFS_LZMA_DEFAULT
	lzma_error = 0;
	res = unlzma(stream->input, length, NULL, NULL, stream->output, NULL,
							error);
	if (res || lzma_error)
		goto failed;
#else
	res = unlzma_improved(stream->input, length, stream->output);
	if (res) {
		ERROR("error -%d\n", res);
		goto failed;
	}
#endif

	/* uncompressed size is stored in the LZMA header (5 byte offset) */
	res = bytes = (int)get_unaligned_le32(stream->input + 5);
	data = squashfs_first_page(output);
	buff = stream->output;
	while (data) {
		if (bytes <= PAGE_SIZE) {
			memcpy(data, buff, bytes);
			break;
		}
		memcpy(data, buff, PAGE_SIZE);
		buff += PAGE_SIZE;
		bytes -= PAGE_SIZE;
		data = squashfs_next_page(output);
	}
	squashfs_finish_page(output);

	return res;

failed:
	ERROR("lzma decompression failed, data probably corrupt\n");
	return -EIO;
}

const struct squashfs_decompressor squashfs_lzma_comp_ops = {
	.init = lzma_init,
	.free = lzma_free,
	.decompress = lzma_uncompress,
	.id = LZMA_COMPRESSION,
	.name = "lzma",
	.supported = 1
};
